{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BERT Classifier Fine Tuning using IMDB Movie Reviews\n",
    "\n",
    "This notebook uses the BERT classifier training scripts from the model zoo to\n",
    "do fine tuning. The [IMDB dataset](https://ai.stanford.edu/~amaas/data/sentiment/)\n",
    "is used to do sentiment analysis on movie reviews. For more information on the large movie\n",
    "review dataset, please see the [ACL 2011 paper](https://aclanthology.org/P11-1015/).\n",
    "\n",
    "Steps:\n",
    "1. [Download the IMDB dataset](#Download-the-IMDB-dataset)\n",
    "2. [Convert the dataset to .tsv files](#Convert-the-dataset-to-.tsv-files)\n",
    "3. [Download the BERT pretrained model checkpoints](#Download-the-BERT-pretrained-model-checkpoints)\n",
    "3. [BERT classifier fine tuning using the IMDB tsv files](#BERT-classifier-fine-tuning-using-the-IMDB-tsv-files)\n",
    "4. [Export the saved model](#Export-the-saved-model)\n",
    "5. [Load the saved model and make predictions](#Load-the-saved-model-and-make-predictions)\n",
    "6. [Quantization using the Intel® Neural Compressor](#Quantization)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and initialization\n",
    "\n",
    "Before starting there are a few dependencies to install and variables to setup. This notebook assumes that TensorFlow has already been installed. There are environment variables to define the location where the [Model Zoo for Intel® Architecture](https://github.com/intelai/models) has been cloned and directories for checkpoint files, the dataset, and output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "if [[ -n \"$(which apt-get)\" ]]; then\n",
    "  apt-get -qq update && apt-get -qq install -y unzip wget\n",
    "elif [[ -n \"$(which apt-get)\" ]]; then\n",
    "  yum install -y unzip wget\n",
    "else\n",
    "  echo \"Please install wget and unzip, or manually download and extract the BERT base checkpoints to ${CHECKPOINT_DIR}\"\n",
    "fi\n",
    "pip install --upgrade -q pip && pip install -q pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "import os\n",
    "import random\n",
    "import sys\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell has parameters that you may want to change."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define path to the model zoo directory, if the env var is not set\n",
    "if \"MODEL_ZOO_DIR\" not in os.environ:\n",
    "    os.environ[\"MODEL_ZOO_DIR\"] = os.path.join(os.environ[\"HOME\"], \"intelai_models\")\n",
    "    \n",
    "# Define the directory where BERT base checkpoints will be downloaded, if the env var is not set\n",
    "if \"CHECKPOINT_DIR\" not in os.environ:\n",
    "    os.environ[\"CHECKPOINT_DIR\"] = os.path.join(os.environ[\"HOME\"], \"bert_base_checkpoints\")\n",
    "\n",
    "# Define the directory where the IMDB dataset will be downloaded/extracted, if the env var is not set\n",
    "if \"DATASET_DIR\" not in os.environ:\n",
    "    os.environ[\"DATASET_DIR\"] = os.path.join(os.environ[\"HOME\"], \"imdb_dataset\")\n",
    "\n",
    "# Define the directory where output logs and checkpoints will be written, if the env var is not set\n",
    "if \"OUTPUT_DIR\" not in os.environ:\n",
    "    os.environ[\"OUTPUT_DIR\"] = os.path.join(os.environ[\"HOME\"], \"bert_classifier_output\")\n",
    "\n",
    "# The number of reviews to use for training (the IMDB dataset has 25,000 training reviews)\n",
    "# Increasing the number of examples used can improve accuracy, but will also increase training time\n",
    "num_train_reviews = 3000\n",
    "\n",
    "# The number of reviews to use for evaluation (the IMDB dataset has 25,000 training reviews)\n",
    "num_test_reviews = 3000\n",
    "\n",
    "# The number of training epochs to run\n",
    "num_train_epochs = 1\n",
    "\n",
    "# Training batch size\n",
    "batch_size = 32\n",
    "\n",
    "# Learning rate\n",
    "learning_rate = \"3e-5\"\n",
    "\n",
    "# Maximum total input sequence length after WordPiece tokenization (longer sequences will be truncated)\n",
    "max_seq_length = 128\n",
    "\n",
    "# Precision (fp32 or bfloat16)\n",
    "# If the model is trained using FP32, the neural compressor can be used to quantize the model after training\n",
    "precision = \"fp32\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Location with the model zoo repo code\n",
    "model_zoo_dir = os.environ[\"MODEL_ZOO_DIR\"]\n",
    "print(\"MODEL_ZOO_DIR:\", model_zoo_dir)\n",
    "if not os.path.isdir(model_zoo_dir):\n",
    "    raise ValueError(\"The model zoo directory ({}) does not exist. This directory should \"\n",
    "                     \"have a clone of the model zoo repo.\".format(model_zoo_dir))\n",
    "sys.path.append(os.path.join(model_zoo_dir, \"models\", \"language_modeling\", \"tensorflow\",\n",
    "                             \"bert_large\", \"training\", precision))\n",
    "\n",
    "# Location where the bert base uncased checkpoints will be downloaded and extracted\n",
    "bert_base_checkpoint_dir = os.environ[\"CHECKPOINT_DIR\"]\n",
    "print(\"CHECKPOINT_DIR:\", bert_base_checkpoint_dir)\n",
    "if not os.path.isdir(bert_base_checkpoint_dir):\n",
    "    os.makedirs(bert_base_checkpoint_dir)\n",
    "\n",
    "# Location where the dataset will be downloaded\n",
    "dataset_dir = os.environ[\"DATASET_DIR\"]\n",
    "print(\"DATASET_DIR:\", dataset_dir)\n",
    "if not os.path.isdir(dataset_dir):\n",
    "    os.makedirs(dataset_dir)\n",
    "\n",
    "# Output directory for logs and checkpoints generated during training\n",
    "output_dir = os.environ[\"OUTPUT_DIR\"]\n",
    "print(\"OUTPUT_DIR:\", output_dir)\n",
    "if not os.path.isdir(output_dir):\n",
    "    os.makedirs(output_dir)\n",
    "\n",
    "if len(os.listdir(output_dir)) != 0:\n",
    "    print(\"\\nWARNING: The OUTPUT_DIR is not empty.\")\n",
    "    print(\"To prevent previously generated checkpoint files from getting picked up, you may want \"\n",
    "          \"to provide an empty output directory.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the IMDB dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next section downloads and extracts the IMDB movie review dataset. This may take a few minutes, depending on your network speed and hardware. If the `aclImdb` folder is already found in the `DATASET_DIR`, the download will be skipped."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "if os.path.isdir(os.path.join(dataset_dir, \"aclImdb\")):\n",
    "    imdb_dataset_dir = os.path.join(dataset_dir, \"aclImdb\")\n",
    "    print(\"Skipping download, since the imdb dataset folder was already found\")\n",
    "else:\n",
    "    imdb_download_url = \"https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz\"\n",
    "    imdb_file_name = os.path.basename(imdb_download_url)\n",
    "\n",
    "    dataset_file = tf.keras.utils.get_file(imdb_file_name, imdb_download_url, untar=True,\n",
    "                                           cache_dir=dataset_dir, cache_subdir=\"\")\n",
    "    imdb_dataset_dir = os.path.join(os.path.dirname(dataset_file), \"aclImdb\")\n",
    "\n",
    "    if not os.path.isdir(imdb_dataset_dir):\n",
    "        raise RuntimeError(\"The IMDB dataset {} folder could not be found at {}.\")\n",
    "    \n",
    "print(\"IMDB dataset dir:\", imdb_dataset_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert the dataset to .tsv files\n",
    "\n",
    "The BERT training scripts expect the dataset to be in .tsv files. The downloaded dataset has folders for `train` and `test` files, and within those folders there are folders named `neg` and `pos` which have negative and positive movie reviews as `.txt` files.\n",
    "\n",
    "The code below shuffles the list of text files and randomly grabs a positive or negative review for the number of entries that are being using in this example for training and test.\n",
    "\n",
    "We use the IMDB data that was just generated to create `train.tsv`, `dev.tsv`, and `test.tsv`. In these `.tsv` files, column 1 has the label (`0` for negative and `1` for positive) and column 3 has the movie review sentence. Column 0 and 2 are unused."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tsv_dir = os.path.join(dataset_dir, \"imdb_tsv\")\n",
    "if not os.path.isdir(tsv_dir):\n",
    "    os.makedirs(tsv_dir)\n",
    "    \n",
    "for data_folder in [\"train\", \"test\"]:\n",
    "    counts = [0, 0]\n",
    "    file_list = [os.listdir(os.path.join(imdb_dataset_dir, data_folder, \"neg\")),\n",
    "                 os.listdir(os.path.join(imdb_dataset_dir, data_folder, \"pos\"))]\n",
    "    random.shuffle(file_list[0])\n",
    "    random.shuffle(file_list[1])\n",
    "    pos_neg = [\"neg\", \"pos\"]\n",
    "    num_reviews = num_train_reviews if data_folder == \"train\" else num_test_reviews\n",
    "    \n",
    "    # Create a dev.tsv and test.tsv from the test data\n",
    "    file_names = [data_folder]\n",
    "    if data_folder == \"test\":\n",
    "        file_names = [\"test\", \"dev\"]\n",
    "    \n",
    "    for file in file_names:\n",
    "        tsv_file = os.path.join(tsv_dir, \"{}.tsv\".format(file))\n",
    "\n",
    "        with open(tsv_file, \"w\") as out_tsv:\n",
    "            tsv_writer = csv.writer(out_tsv, delimiter='\\t')\n",
    "\n",
    "            for x in range(0, num_reviews):\n",
    "                rand_int = random.randint(0, 1)\n",
    "                label = str(rand_int)\n",
    "                txt_file = os.path.join(imdb_dataset_dir, data_folder, pos_neg[rand_int],\n",
    "                                        file_list[rand_int][counts[rand_int]])\n",
    "                counts[rand_int] += 1\n",
    "\n",
    "                with open(txt_file, \"r\") as data_file:\n",
    "                    tsv_writer.writerow(['', str(rand_int), '', data_file.read()])\n",
    "                \n",
    "        print(\"Wrote {} reviews to {}\".format(num_reviews, tsv_file))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the BERT pretrained model checkpoints\n",
    "\n",
    "Download the `uncased_L-12_H-768_A-12` checkpoints to the `CHECKPOINT_DIR` directory and extract the files. The download is skipped if the file already exists."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "mkdir -p ${CHECKPOINT_DIR}\n",
    "cd ${CHECKPOINT_DIR}\n",
    "if [ ! -f \"uncased_L-12_H-768_A-12.zip\" ]; then\n",
    "    wget https://storage.googleapis.com/bert_models/2018_10_18/uncased_L-12_H-768_A-12.zip\n",
    "else\n",
    "    echo \"Skipping download since uncased_L-12_H-768_A-12.zip already exists\"\n",
    "fi\n",
    "\n",
    "if [ ! -d \"uncased_L-12_H-768_A-12\" ]; then\n",
    "    unzip uncased_L-12_H-768_A-12.zip\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## BERT classifier fine tuning using the IMDB tsv files\n",
    "\n",
    "Run the `launch_benchmark.py` script to run BERT training using the BERT base uncased_L-12_H-768_A-12 checkpoints that were just downloaded as the initial checkpoints. The directory with the IMDB tsv files that were created earlier are used as the dataset directory. Checkpoints that are generated during training will be written to the `OUTPUT_DIR`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set env vars to pass as parameters to the model zoo launch_benchmark.py script\n",
    "os.environ[\"NUM_TRAIN_EPOCHS\"] = str(num_train_epochs)\n",
    "os.environ[\"BATCH_SIZE\"] = str(batch_size)\n",
    "os.environ[\"LEARNING_RATE\"] = str(learning_rate)\n",
    "os.environ[\"MAX_SEQ_LENGTH\"] = str(max_seq_length)\n",
    "os.environ[\"PRECISION\"] = precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "!python ${MODEL_ZOO_DIR}/benchmarks/launch_benchmark.py \\\n",
    "  --model-name=bert_large \\\n",
    "  --precision=${PRECISION} \\\n",
    "  --mode=training \\\n",
    "  --framework=tensorflow \\\n",
    "  --batch-size=${BATCH_SIZE} \\\n",
    "  --output-dir ${OUTPUT_DIR} \\\n",
    "  -- train-option=Classifier \\\n",
    "  task-name=cola \\\n",
    "  do-train=true \\\n",
    "  do-eval=true \\\n",
    "  data-dir=${DATASET_DIR}/imdb_tsv \\\n",
    "  vocab-file=${CHECKPOINT_DIR}/uncased_L-12_H-768_A-12/vocab.txt \\\n",
    "  config-file=${CHECKPOINT_DIR}/uncased_L-12_H-768_A-12/bert_config.json \\\n",
    "  init-checkpoint=${CHECKPOINT_DIR}/uncased_L-12_H-768_A-12/bert_model.ckpt \\\n",
    "  max-seq-length=${MAX_SEQ_LENGTH} \\\n",
    "  learning-rate=${LEARNING_RATE} \\\n",
    "  num-train-epochs=${NUM_TRAIN_EPOCHS} \\\n",
    "  optimized_softmax=True \\\n",
    "  experimental_gelu=False \\\n",
    "  do-lower-case=True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check what files are in the output directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $OUTPUT_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export the saved model\n",
    "\n",
    "Use the checkpoint files that were generated during fine tuning to export a `saved_model.pb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf ${OUTPUT_DIR}/frozen\n",
    "!python ${MODEL_ZOO_DIR}/models/language_modeling/tensorflow/bert_large/inference/export_classifier.py \\\n",
    "  --task_name=cola \\\n",
    "  --output_dir=${OUTPUT_DIR} \\\n",
    "  --precision=${PRECISION} \\\n",
    "  --bert_config_file=${CHECKPOINT_DIR}/uncased_L-12_H-768_A-12/bert_config.json \\\n",
    "  --saved_model \\\n",
    "  --signature_def_name=\"serving_default\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check to make sure that there's a saved model in the `$OUTPUT_DIR/frozen` directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $OUTPUT_DIR/frozen"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the saved model and make predictions\n",
    "\n",
    "Load the saved model from the output directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.compat.v1.enable_resource_variables()\n",
    "reloaded_model = tf.saved_model.load(os.path.join(output_dir, \"frozen\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are using classes and functions from the model zoo BERT training model directory to set the tokenizer and a function to create input examples from the movie review sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tokenization\n",
    "from run_classifier import InputExample\n",
    "from run_classifier import convert_examples_to_features\n",
    "from run_classifier import convert_single_example\n",
    "\n",
    "tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.WARN)\n",
    "vocab_file = os.path.join(bert_base_checkpoint_dir, \"uncased_L-12_H-768_A-12\", \"vocab.txt\")\n",
    "tokenizer = tokenization.FullTokenizer(vocab_file=vocab_file, do_lower_case=True)\n",
    "        \n",
    "# Function to create input examples from movie review sentences\n",
    "def create_examples(sentences):\n",
    "    examples = []\n",
    "    \n",
    "    for (i, sentence) in enumerate(sentences):\n",
    "        # Pass the tokenized sentence as the text_a arg. Just always pass label as 0 since we are predicting\n",
    "        text_a = tokenization.convert_to_unicode(sentence)\n",
    "        examples.append(InputExample(guid=i, text_a=text_a, text_b=None, label=\"0\"))\n",
    "        \n",
    "    return examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we setup a list of movie review sentences to use for predicting. We convert these sentences to input examples, then convert the input examples to features. We send the features as a batch to the saved model's evaluation function and get back the prediction results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of sample movie review sentences\n",
    "movie_reviews = [\n",
    "    \"The movie was fantastic\",\n",
    "    \"The worst movie ever!\",\n",
    "    \"Captivating and creative\",\n",
    "    \"Meh\",\n",
    "    \"I'd rather have a cat claw out my eyes than watch that again\",\n",
    "    \"Full of action and suspense\",\n",
    "    \"Overall pretty boring\"\n",
    "]\n",
    "\n",
    "label_list=[\"0\", \"1\"]\n",
    "labels=[\"Negative\", \"Positive\"]\n",
    "num_examples = len(movie_reviews)\n",
    "\n",
    "# Convert the movie review sentences to examples and then convert the examples to features\n",
    "input_examples = create_examples(movie_reviews)\n",
    "features = convert_examples_to_features(input_examples, label_list, max_seq_length, tokenizer)\n",
    "\n",
    "# Create lists for all the feature inputs so that we can do the prediction as a batch\n",
    "all_input_ids = []\n",
    "all_input_mask = []\n",
    "all_segment_ids = []\n",
    "all_label_ids = []\n",
    "\n",
    "for feature in features:\n",
    "    all_input_ids.append(feature.input_ids)\n",
    "    all_input_mask.append(feature.input_mask)\n",
    "    all_segment_ids.append(feature.segment_ids)\n",
    "    all_label_ids.append(feature.label_id)\n",
    "\n",
    "tf_input_ids = tf.constant(all_input_ids, shape=[num_examples, max_seq_length], dtype=tf.int32)\n",
    "tf_input_mask = tf.constant(all_input_mask, shape=[num_examples, max_seq_length], dtype=tf.int32)\n",
    "tf_segment_ids = tf.constant(all_segment_ids, shape=[num_examples, max_seq_length], dtype=tf.int32)\n",
    "tf_label_ids = tf.constant(all_label_ids, shape=[num_examples], dtype=tf.int32)\n",
    "\n",
    "# Use the serving_default signature from the reloaded model to do predictions on our using the feature lists\n",
    "results = reloaded_model.signatures[\"serving_default\"](\n",
    "    input_ids=tf_input_ids, input_mask=tf_input_mask, segment_ids=tf_segment_ids, label_ids=tf_label_ids)\n",
    "\n",
    "# Print the results\n",
    "for (i, sentence) in enumerate(movie_reviews):\n",
    "    print(\"Movie review:\\t{}\".format(sentence))\n",
    "    print(\"Prediction:\\t{}\\n\".format(labels[results[\"probabilities\"][i].numpy().argmax()]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantization\n",
    "\n",
    "Install the [Intel® Neural Compressor](https://github.com/intel/neural-compressor) to use for tuning and quantization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env DEBIAN_FRONTEND=noninteractive\n",
    "!apt-get update && apt-get install -y  build-essential python3-opencv python3-dev \n",
    "!pip install neural-compressor\n",
    "%env TF_ENABLE_MKL_NATIVE_FORMAT=0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Navigate to the BERT IMDB fine tuning notebook directory, which has the `quantize.py` script and the config yaml file, that will be used by the neural compressor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "cd ${MODEL_ZOO_DIR}/docs/notebooks/bert_classifier_fine_tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "!PYTHONPATH=${MODEL_ZOO_DIR}/models/language_modeling/tensorflow/bert_large/training/${PRECISION} \\\n",
    " python quantize.py \\\n",
    "  --config-yaml=${MODEL_ZOO_DIR}/docs/notebooks/bert_classifier_fine_tuning/imdb.yaml \\\n",
    "  --vocab-file=${CHECKPOINT_DIR}/uncased_L-12_H-768_A-12/vocab.txt \\\n",
    "  --tsv-dir ${DATASET_DIR}/imdb_tsv \\\n",
    "  --max-seq-length ${MAX_SEQ_LENGTH} \\\n",
    "  --output-dir ${OUTPUT_DIR}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After quantization is complete, the `OUTPUT_DIR/quantized` folder has the saved model file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $OUTPUT_DIR/quantized"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll load the the quantized model from the `OUTPUT_DIR/quantized` folder, and run predictions using the same feature lists that were setup earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "quantized_model = tf.saved_model.load(os.path.join(output_dir, \"quantized\"))\n",
    "\n",
    "# Use the serving_default signature from the reloaded model to do predictions on our using the feature lists\n",
    "results = quantized_model.signatures[\"serving_default\"](\n",
    "    input_ids=tf_input_ids, input_mask=tf_input_mask, segment_ids=tf_segment_ids, label_ids=tf_label_ids)\n",
    "\n",
    "# Print the results\n",
    "for (i, sentence) in enumerate(movie_reviews):\n",
    "    print(\"Movie review:\\t{}\".format(sentence))\n",
    "    print(\"Prediction:\\t{}\\n\".format(labels[results[\"loss/Softmax\"][i].numpy().argmax()]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
